from rapidocr_onnxruntime import RapidOCR
import os
from PIL import Image
import glob


def extract_and_map_text(image_path, custom_dict):
    """
    提取单张图片中的文字并与自定义字典进行匹配
    """
    if not os.path.exists(image_path):
        print(f"图片文件不存在: {image_path}")
        return None

    # 初始化OCR并识别文字
    ocr = RapidOCR()
    result, _ = ocr(image_path)

    try:
        # 提取并组合文字
        full_text = ""
        if isinstance(result, list) and len(result) > 0:
            # 遍历所有识别结果并组合
            for item in result:
                if isinstance(item, list) and len(item) > 1 and isinstance(item[1], str):
                    full_text += item[1]

        if not full_text:
            print(f"未识别到文字内容: {image_path}")
            return None

        print(f"\n识别到的文字 ({os.path.basename(image_path)}): {full_text}")

        # 与自定义字典进行匹配
        # 精确匹配
        if full_text in custom_dict:
            matched_value = custom_dict[full_text]
            print(f"精确匹配成功: {full_text} -> {matched_value}")
            return matched_value

        # 模糊匹配（检查是否包含字典中的key）
        for key, value in custom_dict.items():
            if key in full_text:
                print(f"模糊匹配成功: {key} -> {value}")
                return value

        print("未找到匹配的内容")
        return None

    except Exception as e:
        print(f"处理图片出错 ({image_path}): {e}")
        return None


def process_image_folder(folder_path, custom_dict):
    """
    处理文件夹中所有图片
    """
    if not os.path.isdir(folder_path):
        print(f"文件夹不存在: {folder_path}")
        return

    # 支持的图片格式
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.gif', '*.tiff']
    image_paths = []

    # 收集所有图片路径
    for ext in image_extensions:
        image_paths.extend(glob.glob(os.path.join(folder_path, ext)))

    if not image_paths:
        print(f"文件夹中未找到图片文件: {folder_path}")
        return

    # 存储所有结果的字典
    all_results = {}

    # 批量处理图片
    print(f"开始处理文件夹: {folder_path}，共 {len(image_paths)} 张图片")
    for img_path in image_paths:
        img_name = os.path.basename(img_path)
        result = extract_and_map_text(img_path, custom_dict)
        all_results[img_name] = result

    # 输出汇总结果
    print("\n" + "=" * 50)
    print("批量处理结果汇总:")
    for img_name, result in all_results.items():
        if result:
            print(f"{img_name}: {result}")
        else:
            print(f"{img_name}: 无匹配结果")
    print("=" * 50)

    return all_results


if __name__ == "__main__":
    # 传入文件夹路径（原图片路径修改为文件夹）
    folder_path = "./img"  # 处理img文件夹下所有图片

    # DNF职业字典
    dnf_classes = {
        # 鬼剑士（男）
        "剑魂": "极诣·剑魂",
        "奇美拉": "聆风·奇美拉",
        "协战师": "重霄·协战师",
        "驭剑士": "极诣·驭剑士",
        "阿修罗": "极诣·阿修罗",
        "剑影": "极诣·剑影",
        "小魔女": "知源·小魔女",
        "狂战士": "极诣·狂战士",
        "流浪武士": "极诣·流浪武士",
        "缪斯": "聆风·缪斯",
        "旅人": "聆风·旅人"
    }

    # 执行批量处理
    process_image_folder(folder_path, dnf_classes)
